import math
import numpy as np
import time
import os
import random
import csv
import cmath
import matplotlib as mpl
import matplotlib.pyplot as plt
from RA_Cal import ReflectArray_cal_FFT, Para_RA, MASK_rp

# 反射阵布置几何参数
h = 0.11
c = 3e8
f0 = 30e9
lam0 = c/f0
L = 9 * lam0
k_0 = 2 * math.pi * f0 / c
aa = lam0/2
NN = 36  # 一列有几个单元

powers = np.array([8, 4, 2, 1])
def batch_gray_to_decimal(gray_array, bits=4):
    gray_array = gray_array.reshape(-1, bits)  # 按4位分组
    binary_array = gray_array.copy()
    binary_array[:, 1] = binary_array[:, 0] ^ gray_array[:, 1]
    binary_array[:, 2] = binary_array[:, 1] ^ gray_array[:, 2]
    binary_array[:, 3] = binary_array[:, 2] ^ gray_array[:, 3]
    return np.sum(binary_array * powers, axis=1)

class Costfunction_tt():
    def __init__(self, t=200):
        self.TT = t  # 采样点数
        self.RA = ReflectArray_cal_FFT.Reflect_Array(f=f0, R_base=L, h=h,
                            RP_LABA=Para_RA.RP_qf, RP_ele=Para_RA.RP_ele, acell=aa, T=self.TT)
        self.RA.FFT_to_FXT()
        self.Marv = self.RA.arv
        self.ang_s11 = Para_RA.ang_
        self.Mm = len(self.RA.fxt)
        self.Mrmn = self.RA.rmn

        aoo = MASK_rp.UV_Multibeam(0, self.TT)
        self.Mask_n = MASK_rp.mask_n

        self.MASKK = np.array([MASK_rp.UV_Multibeam(0, self.TT), MASK_rp.UV_Multibeam(1, self.TT),
                          MASK_rp.UV_Multibeam(2, self.TT), MASK_rp.UV_Multibeam(3, self.TT),
                          MASK_rp.UV_Sector(self.TT)])

        # print(self.Mask_n/t/t, "Initialization CF...")

        self.n_ele = 0
        self.code_to_broad, self.code_to_broad_semi, self.code_to_broad_quar = [], [], []
        self.xiaoX = np.zeros((self.Mm, self.Mm), dtype=complex)
        self.pos_mn = []
        for i in range(NN):
            for j in range(NN):
                self.xiaoX[i, j] = self.RA.s_ele * self.Marv[i, j] * cmath.exp(1j * (- k_0 * self.Mrmn[i, j]))
                if self.Marv[i, j] > 0:
                    self.code_to_broad.append([i, j])
                    if j < NN / 2:
                        self.code_to_broad_semi.append([i, j])
                        if i < NN / 2:
                            self.code_to_broad_quar.append([i, j])
                    self.n_ele = self.n_ele + 1
                    self.pos_mn.append([-L + i*aa + aa/2, -L + j*aa + aa/2])

        self.code_to_broad, self.code_to_broad_semi, self.code_to_broad_quar = \
            np.array(self.code_to_broad), np.array(self.code_to_broad_semi), np.array(self.code_to_broad_quar)
        self.code_bit = 4
        self.n_code = self.n_ele * self.code_bit
        print("CostFunction Code Length: ", self.n_ele*4, "tt: ", self.Mm) # 单元个数， 方向图范围。

    def CostFunction(self, Casei, test_code):
        daX = np.zeros((self.TT, self.TT), dtype=complex)
        test_code = np.array(test_code, dtype=int)

        code10 = batch_gray_to_decimal(test_code)
        woki = range(len(code10))
        angg = self.ang_s11[code10[woki]]

        if Casei == 1:
            xx, yy = self.code_to_broad[woki, 0], self.code_to_broad[woki, 1]
            xiaoo = self.xiaoX[self.code_to_broad[woki, 0], self.code_to_broad[woki, 1]] * \
                    np.exp(1j * angg)
            daX[xx, yy] = xiaoo
        elif Casei == 0:
            xx, yy = self.code_to_broad_quar[woki, 0], self.code_to_broad_quar[woki, 1]
            xiaoo = self.xiaoX[self.code_to_broad_quar[woki, 0], self.code_to_broad_quar[woki, 1]] * \
                    np.exp(1j * angg)
            daX[xx, yy] = xiaoo
            daX[xx, NN - 1 - yy] = xiaoo
            daX[NN - 1 - xx, yy] = xiaoo
            daX[NN - 1 - xx, NN - 1 - yy] = xiaoo
        elif Casei == 4:
            xx, yy = self.code_to_broad_semi[woki, 0], self.code_to_broad_semi[woki, 1]
            xiaoo = self.xiaoX[self.code_to_broad_semi[woki, 0], self.code_to_broad_semi[woki, 1]] * \
                    np.exp(1j * angg)
            daX[xx, yy] = xiaoo
            daX[xx, NN - 1 - yy] = xiaoo

        self.RA.daX = daX
        self.RA.FFT_to_FXT()
        fxt_test = self.RA.fxt
        costf_U = fxt_test - self.MASKK[Casei, 0]
        costf_U = costf_U[costf_U > 0]
        costf_U = np.sum(costf_U**2)

        costf_L = self.MASKK[Casei, 1] - fxt_test
        costf_L = costf_L[costf_L > 0]
        costf_L = np.sum(costf_L**2)
        # return float(- (costf_L+costf_U) / self.Mask_n)
        # return np.log10(costf_L/self.Mask_n), np.log10(costf_U/self.Mask_n)
        return -np.log10((costf_L + costf_U) / (self.TT * self.TT))

    def CostFunction_View(self, Casei, test_code):

        pmni = np.zeros((NN, NN), dtype=complex)
        daX = np.zeros((self.TT, self.TT), dtype=complex)
        test_code = np.array(test_code, dtype=int)

        st = time.time()
        code10 = batch_gray_to_decimal(test_code)
        print("Coding time: ", time.time() - st)

        woki = range(len(code10))
        angg = self.ang_s11[code10[woki]]

        if Casei == 1:
            xx, yy = self.code_to_broad[woki, 0], self.code_to_broad[woki, 1]
            xiaoo = self.xiaoX[self.code_to_broad[woki, 0], self.code_to_broad[woki, 1]] * \
                    np.exp(1j * angg)
            pmni[xx, yy] = angg
            daX[xx, yy] = xiaoo
        elif Casei == 0:
            xx, yy = self.code_to_broad_quar[woki, 0], self.code_to_broad_quar[woki, 1]
            xiaoo = self.xiaoX[self.code_to_broad_quar[woki, 0], self.code_to_broad_quar[woki, 1]] * \
                    np.exp(1j * angg)
            pmni[xx, yy] = angg
            pmni[xx, NN-1-yy] = angg
            pmni[NN-1-xx, yy] = angg
            pmni[NN-1-xx, NN-1-yy] = angg
            daX[xx, yy] = xiaoo
            daX[xx, NN-1-yy] = xiaoo
            daX[NN-1-xx, yy] = xiaoo
            daX[NN-1-xx, NN-1-yy] = xiaoo
        elif Casei == 4:
            xx, yy = self.code_to_broad_semi[woki, 0], self.code_to_broad_semi[woki, 1]
            xiaoo = self.xiaoX[self.code_to_broad_semi[woki, 0], self.code_to_broad_semi[woki, 1]] * \
                    np.exp(1j * angg)
            pmni[xx, yy] = angg
            pmni[xx, NN-1-yy] = angg
            daX[xx, yy] = xiaoo
            daX[xx, NN-1-yy] = xiaoo

        self.RA.daX = daX
        print("Coding time: ", time.time() - st)


        st = time.time()
        self.RA.FFT_to_FXT()
        print("FFT time: ", time.time() - st)

        self.RA.plot_FFT()
        self.RA.Change_ele(pmni, np.ones((NN, NN)))
        self.RA.plot_ELE(self.RA.pha_mn)

        st = time.time()
        fxt_test = self.RA.fxt
        fxt_test_db = np.log10(fxt_test) * 10
        # with open("case1_theory_22.csv", 'a', encoding='utf-8', newline='') as f:
        #     write = csv.writer(f)
        #     write.writerow(list(fxt_test_db[int(self.TT/2), :].round(5)))

        costf_U = fxt_test - self.MASKK[Casei, 0]
        costf_U = costf_U[costf_U > 0]
        costf_U = np.sum(costf_U ** 2)

        costf_L = self.MASKK[Casei, 1] - fxt_test
        costf_L = costf_L[costf_L > 0]
        costf_L = np.sum(costf_L ** 2)
        print("sum time: ", time.time() - st)
        print("Cost at MASK U:", costf_U)
        print("Cost at MASK L:", costf_L)

        return -np.log10((costf_L + costf_U) / (self.TT * self.TT))

    def numberMAP(self):
        fig, ax = plt.subplots(figsize=(L, L))  # 可以调整画布大小
        # 设置坐标轴范围（xmin, xmax, ymin, ymax）
        ax.set_xlim(-L, L)
        ax.set_ylim(-L, L)
        kk = 0
        for i in range(len(self.pos_mn)):
            # 使用 text() 显示数字
            if self.pos_mn[i][1] < 0: # and self.pos_mn[i][1]<0:
                ax.text( -self.pos_mn[i][0], -self.pos_mn[i][1], str(kk),  # 数字及其位置
            ha='center', va='center', fontsize=12,  color='red',)
                kk=kk+1
        # 移除坐标轴（可选，如果不需要坐标轴）
        # ax.axis("equal")
        ax.axis('off')
        plt.show()

if __name__ == '__main__':
    import csvtest
    Code_ = csvtest.Read_csv_inrows(r"codecc.csv", range(int(3808 / 4)))[:900, :]
    samp = len(Code_)

    delta_CF0 = np.zeros(samp)
    # Costfunction_1000 = Costfunction_tt(100)
    # Costfunction_1000.numberMAP()

    # for i in range(samp):
    #     delta_CF0[i] = Costfunction_1000.CostFunction(0, Code_[i])
    #
    for tti in range(50, 501, 50):
        timett = 0
        Costfunction_tti = Costfunction_tt(tti)
        lll, uuu = np.zeros(samp), np.zeros(samp)
        for i in range(samp):
            # popu_rand = [random.randint(0, 1) for _ in range(952)]
            dzt = time.time()
            lll[i], uuu[i]= Costfunction_tti.CostFunction(0, Code_[i])
            timett += time.time()-dzt

        if tti > 0:
            with open("ll_tt.csv", 'a', encoding='utf-8', newline='') as f:
                write = csv.writer(f)
                # zii = np.log10(np.abs(delta_CF - delta_CF0)/delta_CF0)
                # print(np.median(zii))
                write.writerow(list(lll.round(5)))
            with open("uu_tt.csv", 'a', encoding='utf-8', newline='') as f:
                write = csv.writer(f)
                # zii = np.log10(np.abs(delta_CF - delta_CF0)/delta_CF0)
                # print(np.median(zii))
                write.writerow(list(uuu.round(5)))

        # with open("Time_tt.csv", 'a', encoding='utf-8', newline='') as f:
        #     write = csv.writer(f)
        #     write.writerow([tti, timett.__round__(5)])